import numpy as np
import pandas as pd 
import matplotlib.pyplot as plt 
import torch
import pywt

def plot_wave(src, tgt):
    left = []
    right = []
    for index, t in enumerate(tgt):
        if int(t) == 1:
            left.append(src[index])

    for index, t in enumerate(tgt):
        if int(t) == 2:
            right.append(src[index])
            
    for i in range(10):
        channel = 0
        input1 = input("是否查看")
        if input1 == "q":
            break 
        l1 = plt.plot(left[i][:, channel])
        l2 = plt.plot(right[i][:, channel])
        plt.legend(handles = [l1, l2,], labels = ['左右想象', '右手想象'], loc = 'best')
        plt.show()
 
def test_pywt():
    import numpy as np
    import matplotlib.pyplot as plt
    import pywt
    import pywt.data
    
    # Load image
    original = pywt.data.camera()
    
    # Wavelet transform of image, and plot approximation and details
    titles = ['Approximation', ' Horizontal detail',
              'Vertical detail', 'Diagonal detail']
    coeffs2 = pywt.dwt2(original, 'bior1.3')
    LL, (LH, HL, HH) = coeffs2
    plt.imshow(original)
    plt.colorbar(shrink=0.8)
    fig = plt.figure(figsize=(12, 3))
    # for i, a in enumerate([LL, LH, HL, HH]):
    #     ax = fig.add_subplot(1, 4, i + 1)
    #     ax.imshow(a, interpolation="nearest", cmap=plt.cm.gray)
    #     ax.set_title(titles[i], fontsize=10)
    #     ax.set_xticks([])
    #     ax.set_yticks([])
    
    # fig.tight_layout()
    # plt.show()

    print(original.shape)
    print(LL.shape)
    print(LH.shape)


# test_pywt()

if __name__ == "__main__":
    src = torch.load("./src_data_raw.pkl")
    # coeffs2 = pywt.dwt2(src, 'db2')
    # tgt = torch.load("./tgt_data.pkl")
    # plot_wave(coeffs2[0], tgt)

    singal = np.fft.fft(src)

    print(singal.shape)
    plt.plot(singal[0])
    plt.show()
    




